/*
 *  kmeans.cpp
 *  bgsubtraction
 *
 *  Created by a1gucis on 5/12/09.
 *  Copyright 2009 __MyCompanyName__. All rights reserved.
 *
 */

#include "kmeans.h"

KMeans::KMeans(int K) {
	this->K = K;
	this->iterations = KMEANS_ITERATIONS;
	means = cvCreateMat(K, 3, CV_32FC1);
	srand(time(NULL));
	//cout<<rand()/RAND_MAX<<endl;
	for (int i=0;i<K;i++) {
		cvmSet(means, i, 0, ((float)rand()*W_CH1/(float)RAND_MAX)*2*M_PI);
		cvmSet(means, i, 1, (float)rand()*W_CH2/(float)RAND_MAX);
		cvmSet(means, i, 2, (float)rand()*W_CH3/(float)RAND_MAX);
	}
	covMatrices = new vector<CvMat *>(K);
	Utils::PrintMat(means);
}

void KMeans::calcCovMatrices(IplImage *image) {
	for (int i=0;i<covMatrices->size();i++) {
		CvMat *mat = (*covMatrices)[i];
		if (mat != NULL)
			cvReleaseMat(&mat);
	}
	CvMat *chan1Mat = cvCreateMat(image->height, image->width, CV_32FC1);
	CvMat *chan2Mat = cvCreateMat(image->height, image->width, CV_32FC1);
	CvMat *chan3Mat = cvCreateMat(image->height, image->width, CV_32FC1);
	CvMat *mean1Mat = cvCreateMat(image->height, image->width, CV_32FC1);
	CvMat *mean2Mat = cvCreateMat(image->height, image->width, CV_32FC1);
	CvMat *mean3Mat = cvCreateMat(image->height, image->width, CV_32FC1);
	for (int i=0;i<image->height;i++) {
		for (int j=0;j<image->width;j++) {
			float chan1 = ((float *)(image->imageData + i*image->widthStep))[j*image->nChannels + 0];
			float chan2 = ((float *)(image->imageData + i*image->widthStep))[j*image->nChannels + 1];
			float chan3 = ((float *)(image->imageData + i*image->widthStep))[j*image->nChannels + 2];
			cvmSet(chan1Mat, i, j, chan1);
			cvmSet(chan2Mat, i, j, chan2);
			cvmSet(chan3Mat, i, j, chan3);
		}
	}
	CvMat *distch1Mat = cvCreateMat(image->height, image->width, CV_32FC1);
	CvMat *distch2Mat = cvCreateMat(image->height, image->width, CV_32FC1);
	CvMat *distch3Mat = cvCreateMat(image->height, image->width, CV_32FC1);
	CvMat *tmpMat = cvCreateMat(image->height, image->width, CV_32FC1);
	float scale = image->height*image->width;
	for (int z=0;z<K;z++) {
		CvMat *covarMat = cvCreateMat(3, 3, CV_32FC1);
		cvZero(mean1Mat);
		cvAddS(mean1Mat, cvScalar(cvmGet(means, z, 0)), mean1Mat);
		cvZero(mean2Mat);
		cvAddS(mean2Mat, cvScalar(cvmGet(means, z, 1)), mean2Mat);
		cvZero(mean3Mat);
		cvAddS(mean3Mat, cvScalar(cvmGet(means, z, 2)), mean3Mat);
		
		cvSub(chan1Mat, mean1Mat, distch1Mat);
		cvSub(chan2Mat, mean2Mat, distch2Mat);
		cvSub(chan3Mat, mean3Mat, distch3Mat);
		
		//first covar mat row
		cvMul(distch1Mat, distch1Mat, tmpMat);
		float sum11 = Utils::MatrixSum(tmpMat);
		sum11 /= scale;
		cvmSet(covarMat, 0, 0, sum11);
		
		cvMul(distch1Mat, distch2Mat, tmpMat);
		float sum12 = Utils::MatrixSum(tmpMat);
		sum12 /= scale;
		cvmSet(covarMat, 0, 1, sum12);
		
		cvMul(distch1Mat, distch3Mat, tmpMat);
		float sum13 = Utils::MatrixSum(tmpMat);
		sum13 /= scale;
		cvmSet(covarMat, 0, 2, sum13);
		
		//second covar mat row
		cvmSet(covarMat, 1, 0, sum12);
		
		cvMul(distch2Mat, distch2Mat, tmpMat);
		float sum22 = Utils::MatrixSum(tmpMat);
		sum22 /= scale;
		cvmSet(covarMat, 1, 1, sum22);
		
		cvMul(distch2Mat, distch3Mat, tmpMat);
		float sum23 = Utils::MatrixSum(tmpMat);
		sum23 /= scale;
		cvmSet(covarMat, 1, 2, sum23);
		
		//third covar mat row
		cvmSet(covarMat, 2, 0, sum13);
		cvmSet(covarMat, 2, 1, sum23);
		
		cvMul(distch3Mat, distch3Mat, tmpMat);
		float sum33 = Utils::MatrixSum(tmpMat);
		sum33 /= scale;
		cvmSet(covarMat, 2, 2, sum33);
		
		(*covMatrices)[z] = covarMat;
		//covMatrices->push_back(covarMat);
		
	}
	cvReleaseMat(&tmpMat);
	cvReleaseMat(&mean1Mat);
	cvReleaseMat(&mean2Mat);
	cvReleaseMat(&mean3Mat);
	cvReleaseMat(&distch1Mat);
	cvReleaseMat(&distch2Mat);
	cvReleaseMat(&distch3Mat);
	cvReleaseMat(&chan1Mat);
	cvReleaseMat(&chan2Mat);
	cvReleaseMat(&chan3Mat);
}

int KMeans::minMahalanobisDistance(float chan1, float chan2, float chan3) {
	CvMat *chanMat = cvCreateMat(1, 3, CV_32FC1);
	CvMat *meanMat = cvCreateMat(1, 3, CV_32FC1);
	cvmSet(chanMat, 0, 0, chan1);
	cvmSet(chanMat, 0, 1, chan2);
	cvmSet(chanMat, 0, 2, chan3);
	
	cvmSet(meanMat, 0, 0, cvmGet(means, 0, 0));
	cvmSet(meanMat, 0, 1, cvmGet(means, 0, 1));
	cvmSet(meanMat, 0, 2, cvmGet(means, 0, 2));
	CvMat *rangeMat = cvCreateMat(3, 3, CV_32FC1);
	int minPos = 0;
	cvInvert((*covMatrices)[0], rangeMat, CV_SVD_SYM);
	double minDist = cvMahalanobis(chanMat, meanMat, rangeMat);
	for (int i=1;i<K;i++) {
		cvmSet(meanMat, 0, 0, cvmGet(means, i, 0));
		cvmSet(meanMat, 0, 1, cvmGet(means, i, 1));
		cvmSet(meanMat, 0, 2, cvmGet(means, i, 2));
		cvInvert((*covMatrices)[i], rangeMat, CV_SVD_SYM);
		double dist = cvMahalanobis(chanMat, meanMat, rangeMat);
		if (dist < minDist) {
			minDist = dist;
			minPos = i;
		}
	}
	cvReleaseMat(&meanMat);
	cvReleaseMat(&chanMat);
	cvReleaseMat(&rangeMat);
	return minPos+1;
}
int KMeans::minEuclideanDistance(float chan1, float chan2, float chan3) {
	int minPos = 0;
	float dist1 = chan1 - cvmGet(means, 0, 0);
	//cout<<cvmGet(means, 0, 0)<<endl;
	dist1 *= dist1;
	float dist2 = chan2 - cvmGet(means, 0, 1);
	dist2 *= dist2;
	float dist3 = chan3 - cvmGet(means, 0, 2);
	dist3 *= dist3;
	float minDist = dist1 + dist2 + dist3;
	for (int i=1;i<means->rows;i++) {
		dist1 = chan1 - cvmGet(means, i, 0);
		dist1 *= dist1;
		dist2 = chan2 - cvmGet(means, i, 1);
		dist2 *= dist2;
		dist3 = chan3 - cvmGet(means, i, 2);
		dist3 *= dist3;
		float dist = dist1 + dist2 + dist3;
		if (dist < minDist) {
			minDist = dist;
			minPos = i;
		}
	}
	return minPos+1;
}

void KMeans::recalcMeans(IplImage *image) {
	CvMat *sums = cvCreateMat(means->rows, means->cols, CV_64FC1);
	CvMat *hits = cvCreateMat(1, means->rows, CV_32FC1);
	cvZero(sums);
	cvZero(hits);
	for (int i=0;i<image->height;i++) {
		int *labelsPtr = (int *)(labels->data.ptr + i *labels->step);
		for (int j=0;j<image->width;j++) {
			float chan1 = ((float *)(image->imageData + i*image->widthStep))[j*image->nChannels + 0];
			float chan2 = ((float *)(image->imageData + i*image->widthStep))[j*image->nChannels + 1];
			float chan3 = ((float *)(image->imageData + i*image->widthStep))[j*image->nChannels + 2];
			if (isnan(chan1)){
				cout<<"ERROR, chan1 is nan, row:"<<i<<" col"<<j<<endl;
				exit(-1);
			}
			int label = *labelsPtr++;
		//	if ((label > means->rows) || (label < 0))
		//		cout<<"ERROR: "<<label<<endl;
			double sumChan1 = cvmGet(sums, label-1, 0);
			double sumChan2 = cvmGet(sums, label-1, 1);
			double sumChan3 = cvmGet(sums, label-1, 2);
			//cout<<chan1<<" "<<chan2<<" "<<chan3<<endl;
			if (isnan(sumChan1)) {
				cout<<"ERROR, cvmSet's fault"<<endl;
			}
			sumChan1 += chan1;
			sumChan2 += chan2;
			sumChan3 += chan3;
			if (isnan(sumChan1)) {
				cout<<"After addition fault appears, chan1: "<<chan1<<" row:"<<i<<" col:"<<j<<endl;
				exit(-1);
			}
			cvmSet(sums, label-1, 0, sumChan1);
			cvmSet(sums, label-1, 1, sumChan2);
			cvmSet(sums, label-1, 2, sumChan3);
			//cout<<"here done"<<endl;
			int hit = cvmGet(hits, 0, label-1);
			cvmSet(hits, 0, label-1, hit+1);
		}
	}
	for (int i=0;i<means->rows;i++) {
		double sumChan1 = cvmGet(sums, i, 0);
		double sumChan2 = cvmGet(sums, i, 1);
		double sumChan3 = cvmGet(sums, i, 2);
		float hit = cvmGet(hits, 0, i);
		//cout<<"l"<<i+1<<" hit "<<hit<<"\t";
		if (hit > 0) {
			cvmSet(means, i, 0, (float)sumChan1/hit);
			cvmSet(means, i, 1, (float)sumChan2/hit);
			cvmSet(means, i, 2, (float)sumChan3/hit);
			//cout <<"hit: "<<hit<<" |"<< sumChan1 << " "<<sumChan2<< " "<<sumChan3<<endl;
		} else {
			cout<<"hit 0 detected, old mean values: "<<cvmGet(means, i, 0)<<" "<<cvmGet(means, i, 1)<<" "<<cvmGet(means, i, 2)<<endl;
			cvmSet(means, i, 0, ((float)rand()*W_CH1/(float)RAND_MAX)*2*M_PI);
			cvmSet(means, i, 1, (float)rand()*W_CH2/(float)RAND_MAX);
			cvmSet(means, i, 2, (float)rand()*W_CH3/(float)RAND_MAX);	
			cout<<"new means: "<<cvmGet(means, i, 0)<<" "<<cvmGet(means, i, 1)<<" "<<cvmGet(means, i, 2)<<endl;
		}
	}
	//cout<<endl;
	cvReleaseMat(&sums);
	cvReleaseMat(&hits);
}

void KMeans::predict(IplImage *sourceImg, IplImage *labelsImg) {
	labels = cvCreateMat(sourceImg->height, sourceImg->width, CV_32SC1);
	//const float* ptr = (const float*)(mat->data.ptr + row * mat->step); 
	bool labelsConverge = false;
	int sourceImgArea = sourceImg->height * sourceImg->width;
	int tmpIter = 0;
	for (int z=0;z<iterations && !labelsConverge;z++) {
		int labelsMoved = 0;
		calcCovMatrices(sourceImg);
		//cout<<"
		for (int i=0;i<sourceImg->height;i++) {
			int *ptr = (int *)(labels->data.ptr + i *labels->step);
			for (int j=0;j<sourceImg->width;j++) {
				float chan1 = ((float *)(sourceImg->imageData + i*sourceImg->widthStep))[j*sourceImg->nChannels + 0];
				float chan2 = ((float *)(sourceImg->imageData + i*sourceImg->widthStep))[j*sourceImg->nChannels + 1];
				float chan3 = ((float *)(sourceImg->imageData + i*sourceImg->widthStep))[j*sourceImg->nChannels + 2];
				if ((isnan(chan1)) || (isnan(chan2)) || (isnan(chan3))) {
					cout<<"chan1 is nan"<<endl;
					exit(-1);
				}
				//cout<<chan1<<" "<<chan2<<" "<<chan3<<endl;
				int label = minMahalanobisDistance(chan1, chan2, chan3);
				//int label = minEuclideanDistance(chan1, chan2, chan3);
				//cout<<"here 2"<<endl;
				if (label != *ptr)
					labelsMoved++;
				*ptr++ = label;
				//cout<<chan1<<" "<<chan2<<" "<<chan3<<endl;
				//if (++tmpIter == 1)
				//	exit(-1);
			}
		}
		//cout<<"recalcMeans called"<<endl;
		recalcMeans(sourceImg);
		cout<<"iteration: "<<z<<endl;
		float convergence = (float)labelsMoved/(float)sourceImgArea;
		if (convergence < 0.0001) {
			cout<<"sourceImgArea: "<<sourceImgArea<<endl;
			cout<<"labelsMoved: "<<labelsMoved<<endl;
			labelsConverge = true;
		}
	}
	/*cout<<"MEANS: "<<endl;
	Utils::PrintMat(means);
	cout<<"##################"<<endl;
	cout<<"COVARS: "<<endl;
	for (int i=0;i<covMatrices->size();i++) {
		Utils::PrintMat((*covMatrices)[i]);
		cout<<"##################"<<endl;
	}
	cout<<"INVERT: "<<endl;
	for (int i=0;i<covMatrices->size();i++) {
		CvMat *dst = cvCreateMat(3, 3, CV_32FC1);
		cvInvert((*covMatrices)[i], dst, CV_SVD_SYM);
		Utils::PrintMat(dst);
		//if (i == 0) {
		//	cout<<"Multiplication: "<<endl;
		//	CvMat *dst2 = cvCreateMat(3, 3, CV_32FC1);
		//	cvMatMul((*covMatrices)[i], dst, dst2);
		//	Utils::PrintMat(dst2);
		//}
		cvReleaseMat(&dst);
	}
	*/
	uchar *data = (uchar *)labelsImg->imageData;
	for (int i=0;i<labelsImg->height;i++) {
		int *labelsPtr = (int *)(labels->data.ptr + i *labels->step);
		for (int j=0;j<labelsImg->width;j++) {
			data[i*labelsImg->widthStep+j*labelsImg->nChannels] = *labelsPtr++;
		}
	}
	cvReleaseMat(&labels);
}